
import byzerllm
from typing import List, Dict, Optional, Literal, Any, Union
from pydantic import BaseModel
from loguru import logger
import threading
from concurrent.futures import ThreadPoolExecutor
import os
from datetime import datetime
from .types import Message, ReplyDecision
from .utils import GroupUtils,GroupMemberResponse
class AgentHub:
    """管理所有Agent实例的全局中心"""
    _instance = None
    agents: Dict[str, 'Agent'] = {}

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    @classmethod
    def register_agent(cls, agent: 'Agent'):
        """注册一个新的Agent"""
        if agent.name in cls.agents:
            raise ValueError(f"Agent with name {agent.name} already exists")
        cls.agents[agent.name] = agent

    @classmethod
    def get_agent(cls, name: str) -> Optional['Agent']:
        """根据名称获取Agent"""
        return cls.agents.get(name)

    @classmethod
    def list_agents(cls) -> List[str]:
        """获取所有Agent名称"""
        return list(cls.agents.keys())

class GroupHub:
    """管理所有Group实例的全局中心"""
    _instance = None
    groups: Dict[str, 'Group'] = {}
    _groups_lock = threading.RLock()

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    @classmethod
    def register_group(cls, group: 'Group'):
        """注册一个新的Group"""
        with cls._groups_lock:
            if group.name in cls.groups:
                raise ValueError(f"Group with name {group.name} already exists")
            cls.groups[group.name] = group

    @classmethod
    def get_group(cls, name: str) -> Optional['Group']:
        """根据名称获取Group"""
        with cls._groups_lock:
            return cls.groups.get(name)

    @classmethod
    def list_groups(cls) -> List[str]:
        """获取所有Group名称"""
        with cls._groups_lock:
            return list(cls.groups.keys())



class Group:
    # 类级线程池（所有 Group 实例共享）
    _executor = ThreadPoolExecutor(max_workers=min(32, (os.cpu_count() or 1) * 4 + 4))
    
    def __init__(self, name: str):
        self.name = name
        self.members: List[Agent] = []
        self.history: List[Message] = []
        # 成员列表和消息历史的读写锁
        self._members_lock = threading.RLock()
        self._history_lock = threading.RLock()
        # 自动注册到GroupHub
        GroupHub.register_group(self)
        
    def add_member(self, agent: 'Agent'):
        with self._members_lock:
            if agent not in self.members:
                self.members.append(agent)
                with agent._group_lock:  # 访问 Agent 的锁
                    agent.joined_groups[self.name] = self
    
    def _safe_send_message(self, member: 'Agent', message: Message, print_conversation: bool):
        try:                        
            # 调用 Agent 的线程安全接收方法
            member.threadsafe_receive(message,print_conversation=print_conversation)
        except Exception as e:
            import traceback
            traceback.print_exc()
            logger.error(f"Message delivery failed to {member.name}: {str(e)}")

    def _handle_send_result(self, future):
        try:
            future.result()
        except Exception as e:            
            logger.error(f"Message delivery error: {str(e)}")
            
    def broadcast(self, message: Message, print_conversation: bool = False):
        # 消息历史写入需要同步
        with self._history_lock:
            self.history.append(message)
        
        # 使用线程池并行发送
        futures = []
        with self._members_lock:  # 保护成员列表遍历
            current_members = self.members.copy()
        
        for member in current_members:
            if member.name == message.sender:
                continue
                
            futures.append(
                self._executor.submit(
                    self._safe_send_message,
                    member, message, print_conversation
                )
            )

        # 异步等待所有任务完成（可选）
        for future in futures:
            future.add_done_callback(self._handle_send_result)

        if print_conversation:
            with self._history_lock:  # 同步打印
                print(f"[Group Broadcast End] {message.sender} finished broadcasting to {self.name} group")                

class Agent:
    def __init__(self, 
                 name: str, 
                 llm: Union[byzerllm.ByzerLLM, byzerllm.SimpleByzerLLM]):
        self.name = name
        self.llm = llm
        self.joined_groups: Dict[str, Group] = {}
        self.private_chats: Dict[str, List[Message]] = {}
        self.system_prompt = ""
        self.refuse_reply_reason = ""
        self._group_lock = threading.RLock()  # 保护 joined_groups
        self._chat_lock = threading.RLock()   # 保护 private_chats
        # 自动注册到AgentHub
        AgentHub.register_agent(self)

    def who_am_i(self, role: str) -> 'Agent':
        self.system_prompt = role
        return self

    def when_to_refuse_reply(self, reason: str) -> 'Agent':
        self.refuse_reply_reason = reason
        return self
    
    def introduce_myself(self) -> str:        
        return self.system_prompt

        
    def join_group(self, group: Group) -> 'GroupMembership':
        if group.name not in self.joined_groups:
            self.joined_groups[group.name] = group            
            group.add_member(self)            
        return GroupMembership(self, group)    
    
    def talk_to_group(self, group: Group, content: str, mentions: List['Agent'] = [], print_conversation: bool = False):
        message = Message(
            sender=self.name,
            content=content,
            is_group=True,
            group_name=group.name,
            mentions=[m.name for m in mentions]
        )
        group.broadcast(message, print_conversation)
        return self

    def choose_group(self,content:str)->List[GroupMemberResponse]:
        group_utils = GroupUtils(self.llm)
        v = group_utils.auto_select_group(content,self.joined_groups.values())
        return v
            
    def talk_to(self, other: Union['Agent', Group], content: str, mentions: List['Agent'] = [], print_conversation: bool = False):
        if isinstance(other, Group):
            return self.talk_to_group(other, content, mentions, print_conversation)
        
        message = Message(
            sender=self.name,
            content=content,
            is_group=False,
            mentions=[m.name for m in mentions]
        )
        if print_conversation:
            print(f"[Private Chat] {self.name} -> {other.name}: {content}")
            
        # 存储双向对话记录
        self._add_private_message(other.name, message)
        other._add_private_message(self.name, message)                
        
        response = other.generate_reply.with_llm(self.llm).with_return_type(ReplyDecision).run(message)                                
        
        if print_conversation:
            print(f">>> {other.name} reply to {self.name} with strategy: {response.strategy}, reason: {response.reason}")

        if response.strategy == "ignore":            
            return
        elif response.strategy == "private":            
            mentions = [AgentHub.get_agent(m) for m in response.mentions]             
            other.talk_to(other=self, content=response.content, mentions=mentions, print_conversation=print_conversation)
        elif response.strategy == "broadcast":
            warning_msg = f"invalid strategy broadcast action in private chat:[{self.name}] 广播消息给群组 {other.name}: {response.content}"
            logger.warning(warning_msg)
            if print_conversation:
                print(f"[Private Chat Warning] {warning_msg}")
        return self
            
    
    def _add_private_message(self, other_name: str, message: Message):
        with self._chat_lock:
            if other_name not in self.private_chats:
                self.private_chats[other_name] = []
            self.private_chats[other_name].append(message)
    
    def threadsafe_receive(self, message: Message,print_conversation: bool = False):                
        self.receive_message(message,print_conversation=print_conversation)
        
    def receive_message(self, message: Message,print_conversation: bool = False):
        if message.is_group:
            prefix = f"[Group {message.group_name}]"
            if message.mentions and self.name in message.mentions:
                prefix += " @You"
            print(f"{prefix} {message.sender}: {message.content}")                        
            reply_decision = self.generate_reply.with_llm(self.llm).with_return_type(ReplyDecision).run(message)            
            print(f">>> {self.name} reply to {message.sender} with strategy: {reply_decision.strategy}, reason: {reply_decision.reason}")
            if reply_decision.strategy == "ignore":                
                return
            elif reply_decision.strategy == "private":                
                self.talk_to(other=AgentHub.get_agent(message.sender), content=reply_decision.content,print_conversation=print_conversation)
            elif reply_decision.strategy == "broadcast":
                self.joined_groups[message.group_name].broadcast(Message(
                    sender=self.name,
                    content=reply_decision.content,
                    is_group=True,
                    group_name=message.group_name,
                    mentions=reply_decision.mentions,
                    priority=reply_decision.priority
                ))
        else:
            print(f"[Private] {message.sender}: {message.content}")
            
    @byzerllm.prompt()
    def generate_reply(self, message: Message) -> str:
        """
        你的名字是 {{ name }}
        {% if message.is_group %}
        当前群组是 {{ message.group_name }}
        {% endif %}
        当前时间: {{ time }}
        {% if role %}
        你对自己的描述是: 
        <who_are_you>
        {{ role }}
        </who_are_you>
        {% endif %}

        
        {{ message.sender }} 发送了一条{% if message.is_group %}群组消息{% else %}私聊消息{% endif %}：
        <message>
        {{ message.content }}
        </message>
        
        {% if message.mentions and message.mentions|length > 0 %}
        这条消息的发送者特别 @ 了用户：（{{ message.mentions|join(',') }}）。        
        {% endif %}
                
        {% if message.is_group %}
        群组对话上下文：
        <group_message_history>
        {% for msg in group_message_history %}
        {% if msg.sender == name %}
        <role>你/You</role>: <msg>{{ msg.content }}</msg>
        {% else %}
        <role>{{ msg.sender }}</role>: <msg>{{ msg.content }}</msg>
        {% endif %}
        {% endfor %}
        </group_message_history>
        {% else %}
        私聊对话上下文：
        <private_message_history>
        {% for msg in private_message_history %}
        {% if msg.sender == name %}
        <role>你/You</role>: <msg>{{ msg.content }}</msg>
        {% else %}
        <role>{{ msg.sender }}</role>: <msg>{{ msg.content }}</msg>
        {% endif %}
        {% endfor %}
        </private_message_history>
        {% endif %}
        
        请生成 JSON 格式回复：

        ```json
        {
            "content": "回复内容",
            "strategy": "broadcast|private|ignore",
            "mentions": ["被提及的agent名称"],
            "priority": 优先级 0-100,
            "reason": "选择策略的原因"
        }
        ```
        ignore 表示用户发送的消息可以不会进行回复,private 我们要回复用户，并且只回复给发送信息的用户，broadcast 表示要回复消息，并且回复给群组所有成员。
        *** 注意，阅读上面的所有内容，尤其关注 {% if message.is_group %}群组上下文{% else %}私聊上下文{% endif %}，判断是否使用 ignore 策略结束对话，避免无意义对话。一般在群组对话中，你没有被 @ 就无需回答，直接使用ignore策略结束对话。 ***
        {% if refuse_reply_reason %}
        当满足以下描述时，你应当拒绝回复：
        <when_to_refuse_reply>
        {{ refuse_reply_reason }}
        </when_to_refuse_reply>
        {% endif %}
        请严格按照 JSON 格式输出，不要有任何多余的内容。
        """
        context = {
            "name": self.name,
            "role": self.system_prompt,
            "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "refuse_reply_reason": self.refuse_reply_reason
        }
        if message.is_group:
            group = self.joined_groups[message.group_name]
            context["group_message_history"] = group.history
        else:
            context["private_message_history"] = self.private_chats.get(message.sender, [])
            
        return context

class GroupMembership:
    def __init__(self, agent: Agent, group: Group):
        self.agent = agent
        self.group = group
        
    def talk_to_all(self, content: str):
        message = Message(
            sender=self.agent.name,
            content=content,
            is_group=True,
            group_name=self.group.name
        )
        self.group.broadcast(message)
        
    def talk_to(self, content: str, mentions: List[Agent]):
        message = Message(
            sender=self.agent.name,
            content=f"@{' @'.join([a.name for a in mentions])} {content}",
            is_group=True,
            group_name=self.group.name,
            mentions=[a.name for a in mentions]
        )
        self.group.broadcast(message)